package midend.IR_Optimizer;

import midend.Dom_Analysis;
import midend.IR_Instruction.*;
import midend.IR_Type.IR_Int32;
import midend.IR_Type.IR_Float;
import midend.IR_Type.IR_Type;
import midend.IR_Value.*;

import java.util.*;

public class Mem_To_Reg implements IR_Opt{

    // 记录可以被优化掉的alloc指令
    Map<String, IR_Value> allocaVars = new HashMap<>();

    // 记录alloca指令申请的变量在哪些基本块中被定义
    Map<IR_Value, Set<IR_BaseBlock>> varDefBlock = new HashMap<>();

    // 各基本块前驱
    Map<IR_BaseBlock, List<IR_BaseBlock>> predecessors;

    // GraphColoring_RegisterAllocator graphColoring_registerAllocator = new GraphColoring_RegisterAllocator();

    @Override
    public void Optimize(IR_FuncBlock funcBlock){
        /*if (funcBlock.get_bb_List().size() > 2000){
            return;
        }*/
        // 计算支配边界
        Dom_Analysis dom_analysis = new Dom_Analysis(funcBlock);
        Map<IR_BaseBlock, Set<IR_BaseBlock>> dom_frontiers= dom_analysis.getDominanceFrontiers();
        predecessors = dom_analysis.calcuPredecessors(funcBlock);
        for(IR_BaseBlock bb : funcBlock.get_bb_List()) {
            List<Instruction> instructions = bb.get_Instruction_List();
            Iterator<Instruction> iterator = instructions.iterator();
            while(iterator.hasNext()) {
                Instruction instruction = iterator.next();
                if (instruction instanceof Branch_Instruction || instruction instanceof Ret_Instruction) {
                    while (iterator.hasNext()) {
                        iterator.next();
                        iterator.remove();
                    }
                    break;
                }
                else if (instruction instanceof Alloc_Instrction) {
                    if ((((Alloc_Instrction) instruction).get_op_Type().get_TypeName().equals("i32*")) ||
                            (((Alloc_Instrction) instruction).get_op_Type().get_TypeName().equals("float*"))) {
                        IR_Value op = instruction.get_op().get(0);
                        varDefBlock.put(op, new HashSet<>());
                        allocaVars.put(((Alloc_Instrction) instruction).get_op_Name(), op);
                    }
                }
            }
        }

        // 记录alloc指令的定义点
        for (IR_BaseBlock bb : funcBlock.get_bb_List()) {
            for (Instruction instruction : bb.get_Instruction_List()) {
                if (instruction instanceof Store_Instruction) {
                    IR_Value ptr = instruction.get_op().get(2);
                    if(allocaVars.containsKey(ptr.get_ValueName())) {
                        varDefBlock.get(allocaVars.get(ptr.get_ValueName())).add(bb);
                    }
                }
            }
        }

        // 插入phi节点
        Map<Phi_Instruction, IR_Value> addedPhis = insertPhiNodes(dom_frontiers);

        // 更新所有的定义和引用
        //long varRenameStart = System.currentTimeMillis();
        var_Rename(funcBlock, addedPhis, dom_analysis);
        //long varRenameEnd = System.currentTimeMillis();
        //System.out.println("Variable renaming time: " + (varRenameEnd - varRenameStart) + "ms");

        remove_Unused_Phi(funcBlock);
    }

    private Map<Phi_Instruction, IR_Value> insertPhiNodes(Map<IR_BaseBlock, Set<IR_BaseBlock>> domFrontiers) {
        //记录新插入的 phi 指令及其对应的变量
        Map<Phi_Instruction, IR_Value> newPhis = new HashMap<>();
        //记录已插入过phi指令的基本块
        Set<IR_BaseBlock> visited_block = new HashSet<>();

        for (IR_Value var : varDefBlock.keySet()) {
            // 初始化工作队列为变量的定义点所在的所有基本块
            Queue<IR_BaseBlock> workList = new LinkedList<>(varDefBlock.get(var));
            visited_block.clear();

            while (!workList.isEmpty()) {
                IR_BaseBlock block = workList.poll();
                Set<IR_BaseBlock> frontiers = domFrontiers.get(block);
                if (frontiers == null) {
                    continue;
                }
                for (IR_BaseBlock df : frontiers) {
                    // 如果支配边界的基本块未被处理
                    if (visited_block.add(df)) {
                        IR_Type valueType = var.get_Type().get_TypeName().equals("i32*") ? IR_Int32.IR_Int32() : IR_Float.IR_Float();
                        IR_Value resRegister = new IR_VirtualReg_Value("phi", valueType); // 这里的phi_x只是变量名
                        Phi_Instruction phi = new Phi_Instruction(Collections.singletonList(resRegister), df); // 这一步工作只是插入phi指令，还没有对phi指令的数据流进行分析
                        df.add_InstructionList_at_Start(phi);
                        newPhis.put(phi, var);
                        if (!varDefBlock.get(var).contains(df)) {
                            workList.add(df);
                        }
                    }
                }
            }
        }
        return newPhis;
    }


    private void var_Rename(IR_FuncBlock functionBlock, Map<Phi_Instruction, IR_Value> newPhis, Dom_Analysis dom_analysis) {
        //记录已经访问过的基本块
        Set<IR_BaseBlock> visited = new HashSet<>();
        //记录待处理的基本块及其对应的进入值映射
        Queue<Map.Entry<IR_BaseBlock, Map<IR_Value, IR_Value>>> workList = new LinkedList<>(); // Map中有很多个键值对。Map.Entry表示其中的一个键值对
        //初始化
        Map<IR_Value, IR_Value> initialcomingVals = new HashMap<>();
        workList.add(new AbstractMap.SimpleEntry<>(functionBlock.get_bb_List().get(0), initialcomingVals));

        while (!workList.isEmpty()) {
            Map.Entry<IR_BaseBlock, Map<IR_Value, IR_Value>> entry = workList.poll();
            IR_BaseBlock bb = entry.getKey();
            Map<IR_Value, IR_Value> comingVals = entry.getValue();
            //处理过
            if (visited.contains(bb)) continue;
            visited.add(bb);

            //long t1_start = System.currentTimeMillis();
            Iterator<Instruction> iterator = bb.get_Instruction_List().iterator();
            ArrayList<Instruction> notPhiInstructions = new ArrayList<>();
            for (Instruction notPhi : bb.get_Instruction_List()) {
                if (!(notPhi instanceof Phi_Instruction)) {
                    notPhiInstructions.add(notPhi);
                }
            }


            while (iterator.hasNext()) {
                Instruction inst = iterator.next();

                if (inst instanceof Load_Instruction) {  // 示例：%val = load i32, i32* %x，把所有%val引用处的%val，替换为变量的当前版本：incomingVals
                    IR_Value ptr = inst.get_op().get(1);
                    if (allocaVars.containsKey(ptr.get_ValueName())) {
                        //遍历函数的所有基本块，将用到 load 结果的地方都替换成 comingVals[L]
                        IR_Value res = inst.get_op().get(0);
                        for (IR_BaseBlock bblock : functionBlock.get_bb_List()) {
                            for (Instruction instruction : notPhiInstructions) {
                                List<IR_Value> ops = instruction.get_op();
                                //检查call指令中的参数是否引用load的结果
                                if (instruction instanceof Call_Instruction callInstruction) {
                                    for (int k = 0; k < callInstruction.get_params().size(); k++) {
                                        if (callInstruction.get_params().get(k).equals(res)) {
                                            callInstruction.setParams(k, comingVals.get(ptr));
                                        }
                                    }
                                }
                                // 检查其他指令是否引用load指令的结果
                                for (int k = 0; k < ops.size(); k++) {
                                    if (ops.get(k) == null) {
                                        continue;
                                    }
                                    if (ops.get(k).get_ValueName().equals(res.get_ValueName())) {
                                        if (instruction instanceof Pointer_Instruction && k>=2) {
                                            ((Pointer_Instruction) instruction).set_List(k - 2, comingVals.get(ptr));
                                            instruction.set_Op(k, comingVals.get(ptr));
                                        } else if (instruction instanceof Type_Trans_Instruction) {
                                            // 对类型转换指令进行特殊的处理
                                            if (k == 0) {
                                                ((Type_Trans_Instruction) instruction).set_Res(comingVals.get(ptr));
                                            } else if (k == 1) {

                                                ((Type_Trans_Instruction) instruction).set_Origin(comingVals.get(ptr));
                                            }
                                            instruction.set_Op(k, comingVals.get(ptr));
                                        } else {
                                            instruction.set_Op(k, comingVals.get(ptr));
                                        }
                                    }
                                }
                            }
                        }
                        iterator.remove();
                    }
                } else if (inst instanceof Store_Instruction) {  // <result> = store <ty> <value>, <ty>* <pointer>
                    IR_Value pointer = inst.get_op().get(2);
                    if (allocaVars.containsKey(pointer.get_ValueName())) {
                        comingVals.put(pointer, inst.get_op().get(1)); // 将alloca变量的定义值加入到变量版本中，此后用定义值代替引用处的alloca变量
                        iterator.remove();
                    }
                } else if (inst instanceof Phi_Instruction && newPhis.containsKey(inst)) {
                    IR_Value alloca = newPhis.get(inst);
                    comingVals.put(alloca, inst.get_op().get(0));
                } else if (inst instanceof Alloc_Instrction) {
                    if (allocaVars.containsKey(((Alloc_Instrction) inst).get_op().get(0).get_ValueName())) {
                        iterator.remove();
                    }
                }
            }
            //long t1_end = System.currentTimeMillis();
            //System.out.println("t1 time: " + (t1_end - t1_start) + "ms");

            //更新工作集维护phi指令，phi指令有一个单独的comingValues，里面记录着某一变量能到达该phi指令的所有定值
            // 计算支配边界
            //long t2_start = System.currentTimeMillis();
            //Dom_Analysis dom_analysis = new Dom_Analysis(functionBlock);
            for (IR_BaseBlock succ : dom_analysis.get_Immediate_Successors(bb)) {
                //System.out.println(dom_analysis.get_Immediate_Successors(bb).size());
                workList.add(new AbstractMap.SimpleEntry<>(succ, new HashMap<>(comingVals)));
                for (Instruction phiInst : succ.get_Instruction_List()) {
                    if (phiInst instanceof Phi_Instruction && newPhis.containsKey(phiInst)) {
                        IR_Value alloca = newPhis.get(phiInst);
                        IR_Value irValue = comingVals.get(alloca);
                        //这行伪代码里没有，但不加有空指针
                        if (irValue == null) {
                            if (phiInst.get_op().get(0).get_Type() instanceof IR_Int32) {
                                irValue = new IR_Int_Value(0, IR_Int32.IR_Int32());
                            } else {
                                irValue = new IR_Float_Value(0, IR_Float.IR_Float());
                            }
                        }
                        ((Phi_Instruction) phiInst).set_ComingValue(irValue, bb); // 设置phi指令
                    }
                }
            }
            //long t2_end = System.currentTimeMillis();
            //System.out.println("t2 time: " + (t2_end - t2_start) + "ms");
        }
    }


    private void remove_Unused_Phi(IR_FuncBlock functionBlock) {
        for (IR_BaseBlock bb : functionBlock.get_bb_List()) {
            List<Instruction> instructions = bb.get_Instruction_List();
            Iterator<Instruction> iterator = instructions.iterator();
            // 遍历基本块中的所有指令
            while (iterator.hasNext()) {
                Instruction inst = iterator.next();
                if (inst instanceof Phi_Instruction && !isUsed(inst, functionBlock)) {
                    iterator.remove();
                }
            }
        }
    }

    private boolean isUsed(Instruction inst, IR_FuncBlock functionBlock) {
        for (IR_BaseBlock bb : functionBlock.get_bb_List()) {
            for (Instruction instruction : bb.get_Instruction_List()) {
                IR_Value phiRes = inst.get_op().get(0);
                if (inst.equals(instruction)) {
                    continue;
                }
                if (instruction instanceof Call_Instruction callInstruction) {
                    for (IR_Value param : callInstruction.get_params()) {
                        if (param.equals(phiRes)) {
                            return true;
                        }
                    }
                } else if (instruction instanceof Pointer_Instruction pointerInstruction) {
                    for (IR_Value index : pointerInstruction.get_list()) {
                        if (index.equals(phiRes)) {
                            return true;
                        }
                    }
                } else if (instruction instanceof Phi_Instruction phiInstruction) {
                    for (IR_BaseBlock key : phiInstruction.get_ComingValues().keySet()) {
                        if (phiInstruction.get_ComingValues().get(key).equals(phiRes)) {
                            return true;
                        }
                    }
                }
                for (IR_Value op : instruction.get_op()) {
                    if (phiRes.equals(op)) {
                        return true;
                    }
                }
            }
        }
        return false;
    }

}






